MODULE pfmm_box_builder

   USE fmm_global_consts
   USE fmm_global_types
   USE fmm_stats
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: pfmm_init_box_builder,                          &
             pfmm_free_box_builder,                          &
             pfmm_get_box_paras_at_level,                    &
             pfmm_get_box_qlm_at_level

   ! Pointers to unboxed LHS parameters and RHS moments & parameters
   TYPE(raw_mm_data),  POINTER, SAVE :: RHS_raw_mms
   TYPE(raw_mm_paras), POINTER, SAVE :: LHS_raw_paras(:)
   ! Packed LHS & RHS parameters and RHS moments at all levels
   TYPE(box_mm_data), POINTER, SAVE :: mms_at_lev(:)
   ! Box hierarchy ( 1=>supercell; 2=>box 27 times larger than supercell...)
   INTEGER(INTK), SAVE :: pfmm_max_level
   ! Lattice constants, or PFMM box size 
   REAL(REALK), SAVE :: cell_size(3) = (/1d0,1d0,1d0/)

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_init_box_builder(LHS_paras,RHS_mms,max_level_in,lat_consts)

      IMPLICIT NONE
      TYPE(raw_mm_paras), TARGET, INTENT(INOUT) :: LHS_paras(:)
      TYPE(raw_mm_data),  TARGET, INTENT(INOUT) :: RHS_mms
      INTEGER(INTK),              INTENT(IN)    :: max_level_in
      REAL(REALK),                INTENT(IN)    :: lat_consts(3)

      pfmm_max_level = max_level_in
      cell_size(:) = lat_consts(:)

      ! initialise RHS pointers to raw moments and parameters
      RHS_raw_mms => RHS_mms
      ! initialise LHS pointers to raw parameters
      LHS_raw_paras => LHS_paras(:) 
      ! allocate foundation of hierarchical moment data structure
      CALL pfmm_preallocate_levels
      ! update raw_mm_paras with box and branch data
      CALL pfmm_init_box_paras(LHS_paras,RHS_mms%paras) 

   END SUBROUTINE pfmm_init_box_builder

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_init_box_paras(LHS,RHS)

      USE fmm_box_utils, ONLY: fmm_box, fmm_branch, fmm_grain, fmm_box_centre

      IMPLICIT NONE
      TYPE(raw_mm_paras), INTENT(INOUT) :: LHS(:)
      TYPE(raw_mm_paras), INTENT(INOUT) :: RHS(:)
!      TYPE(fmm_scheme),   INTENT(IN)    :: scheme

      INTEGER(INTK) :: i
!      REAL(REALK)   :: grain, grain_inv

!      grain     = fmm_grain(scheme,deepest_level)
!      grain_inv = one/grain
      
      DO i = 1, SIZE(RHS) 
         RHS(i)%box = 1
         RHS(i)%bra = 1
!FIXME coordinate centre of supercell
         RHS(i)%box_cntr(:) = zero
         RHS(i)%map_up = 0
!         RHS(i)%box = fmm_box(RHS(i)%cntr,grain_inv)
!         RHS(i)%bra = fmm_branch(RHS(i)%ext,grain_inv)
!         RHS(i)%box_cntr = fmm_box_centre(RHS(i)%box,grain)
!         RHS(i)%map_up = 0     ! not needed in PFMM
      END DO

      DO i = 1, SIZE(LHS)
         LHS(i)%box = 1
         LHS(i)%bra = 1
         LHS(i)%box_cntr(:) = zero
         LHS(i)%map_up = 0
      END DO

   END SUBROUTINE pfmm_init_box_paras

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_build_paras_at_level(level,scheme)

      USE fmm_box_packer, ONLY: fmm_init_pkd_paras

      IMPLICIT NONE
      TYPE(fmm_scheme), INTENT(IN) :: scheme
      INTEGER(INTK),    INTENT(IN) :: level

      TYPE(raw_mm_paras), POINTER :: ptr(:)

      IF ((level < 1) .OR. (level > pfmm_max_level)) THEN
         CALL fmm_quit ('pfmm cannot iterate paras to this level!')
      END IF

      ! build packed parameters at deepest level if not already done
      IF (.NOT.ASSOCIATED(mms_at_lev(1)%RHS_paras)) THEN
         ptr => RHS_raw_mms%paras(:)
         CALL fmm_init_pkd_paras(1,scheme,ptr,mms_at_lev(1)%RHS_paras)
      END IF
      IF (.NOT.ASSOCIATED(mms_at_lev(1)%LHS_paras)) THEN
         ptr => LHS_raw_paras(:)
         CALL fmm_init_pkd_paras(1,scheme,ptr,mms_at_lev(1)%LHS_paras)
      END IF

      IF (level > 1) THEN
         ! iterate paras through the box hierarchy to the reqd level
         ! (note: all boxes same for simple pfmm!)
         CALL pfmm_copy_paras_to_level(level,scheme)
      END IF

   END SUBROUTINE pfmm_build_paras_at_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_copy_paras_to_level(level,scheme)

      IMPLICIT NONE
      INTEGER(INTK),    INTENT(IN) :: level
      TYPE(fmm_scheme), INTENT(IN) :: scheme

      INTEGER(INTK) :: l, ndim
      
      ndim = SIZE(mms_at_lev(1)%RHS_paras)
      DO l = 2, level
         ALLOCATE (mms_at_lev(l)%RHS_paras(ndim))
         ALLOCATE (mms_at_lev(l)%LHS_paras(ndim))
         mms_at_lev(l)%RHS_paras(:) = mms_at_lev(1)%RHS_paras(:)
         mms_at_lev(l)%LHS_paras(:) = mms_at_lev(1)%LHS_paras(:)
      END DO

   END SUBROUTINE pfmm_copy_paras_to_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_build_qlm_at_level(level,scheme,memory)

      USE fmm_W_pair_builder, ONLY: fmm_translate_raw_moments

      IMPLICIT NONE
      TYPE(fmm_scheme), INTENT(IN) :: scheme
      INTEGER(INTK),    INTENT(IN) :: level
      CHARACTER(4),     INTENT(IN) :: memory

      TYPE(raw_mm_data), POINTER :: ptr1
      TYPE(box_mm_data), POINTER :: ptr2
      INTEGER(INTK) :: mms_dim
      LOGICAL :: flag

      IF ((level < 1) .OR. (level > pfmm_max_level)) THEN
         CALL fmm_quit ('pfmm cannot iterate boxed moments to this level!')
      END IF

      ! build boxed moments at the deepest level if not already done
      IF (.NOT.ASSOCIATED(mms_at_lev(1)%qlm_W)) THEN
         mms_dim = SIZE(mms_at_lev(1)%RHS_paras)
         CALL allocate_lm_at_level(1,mms_dim,scheme%trans_LMAX)
         IF (.NOT.ASSOCIATED(RHS_raw_mms))    &
               CALL fmm_quit('mm_box_builder not correctly initialised!')
         ptr1 => RHS_raw_mms
         ptr2 => mms_at_lev(1)
         CALL fmm_translate_raw_moments(scheme,ptr1,ptr2)
!FIXME -- Even some small noise in the total charge per supercell
!         can quickly blow up when considering an "infinite" far-field.
!         So we explicitly set to exactly 0 to avoid this.
ptr2%qlm_W(1,1) = zero
ptr2%qlm_T(1,1) = zero
! Also set dipole to zero manually
ptr2%qlm_W(2:4,1) = zero
ptr2%qlm_T(2:4,1) = zero
      END IF

      IF (level > 1) THEN
         ! iterate RHS moments through the box hierarchy to the reqd level
         CALL pfmm_iterate_qlm_to_level(level,scheme,memory)
      END IF

      flag = .FALSE.
      IF ( SIZE(mms_at_lev(level)%RHS_paras) > 1 ) flag = .TRUE.
      IF ( SIZE(mms_at_lev(level)%LHS_paras) > 1 ) flag = .TRUE.
      IF (flag) CALL fmm_quit('pfmm too many boxed paras in supercell!')

   END SUBROUTINE pfmm_build_qlm_at_level

!-------------------------------------------------------------------------------

   RECURSIVE SUBROUTINE pfmm_iterate_qlm_to_level(level,scheme,memory)

      IMPLICIT NONE
      INTEGER(INTK),    INTENT(IN) :: level
      TYPE(fmm_scheme), INTENT(IN) :: scheme
      CHARACTER(4),     INTENT(IN) :: memory

      INTEGER(INTK) :: l, mms_dim, l_down
      TYPE(box_mm_data), POINTER :: ptr1, ptr2

      l_down = level-1
      ! note we always build qlm_W first, then scale to give qlm_T
      IF (.NOT.ASSOCIATED(mms_at_lev(l_down)%qlm_W)) THEN
         ! must have qlm at deeper levels before higher levels
         CALL pfmm_iterate_qlm_to_level(l_down,scheme,memory)
      END IF

      ! must build boxed paras before boxed moments
      IF (.NOT.ASSOCIATED(mms_at_lev(level)%RHS_paras)) THEN
         CALL pfmm_build_paras_at_level(level,scheme)
      END IF

      CALL pfmm_combine_periodic_moments(level,scheme)

   END SUBROUTINE pfmm_iterate_qlm_to_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_combine_periodic_moments(level,scheme)

      USE fmm_W_pair_builder, ONLY: fmm_translate_boxed_moments

      IMPLICIT NONE
      INTEGER(INTK),    INTENT(IN) :: level
      TYPE(fmm_scheme), INTENT(IN) :: scheme

      TYPE(box_mm_data), POINTER :: old_ptr, new_ptr
      INTEGER(INTK) :: mms_dim, i,j,k
      REAL(REALK) :: cntr(3), shift(3)

      mms_dim = SIZE(mms_at_lev(level)%RHS_paras)
      CALL allocate_lm_at_level(level,mms_dim,scheme%trans_LMAX)

      old_ptr => mms_at_lev(level-1)
      new_ptr => mms_at_lev(level)

      ! Initialize old to new parameters
      old_ptr%RHS_paras(1)%id = 1
      old_ptr%RHS_paras(1)%map_up = 1
      old_ptr%RHS_paras(1)%cntr(:) = zero
      old_ptr%RHS_paras(1)%cntr_up(:) = zero

      shift = cell_size * (3**(level-2))
      DO i = -1, 1
         DO j = -1, 1
            DO k = -1, 1
               cntr(1) = i*shift(1)
               cntr(2) = j*shift(2)
               cntr(3) = k*shift(3)
               old_ptr%RHS_paras(1)%cntr(:) = cntr(:)
               CALL fmm_translate_boxed_moments(scheme,old_ptr,new_ptr)
            END DO
         END DO
      END DO

   END SUBROUTINE pfmm_combine_periodic_moments

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_get_box_paras_at_level(l,scheme,box_paras,side)

      IMPLICIT NONE
      INTEGER(INTK),      INTENT(IN) :: l
      TYPE(fmm_scheme),   INTENT(IN) :: scheme
      TYPE(box_mm_paras), POINTER    :: box_paras(:)
      CHARACTER(3),       INTENT(IN) :: side 

      IF (.NOT.ASSOCIATED(mms_at_lev)) STOP 'mms_at_lev should be allocated!'

      SELECT CASE (side)
      CASE ('RHS')
         IF (.NOT.ASSOCIATED(mms_at_lev(l)%RHS_paras)) THEN
            ! RHS paras not available at this level yet. So make them!
            CALL pfmm_build_paras_at_level(l,scheme)
         END IF
         box_paras => mms_at_lev(l)%RHS_paras(:)
         stat_RHS_boxes(l) = SIZE(box_paras)
      CASE ('LHS')
         IF (.NOT.ASSOCIATED(mms_at_lev(l)%LHS_paras)) THEN
            ! LHS paras not available at this level yet. So make them!
            CALL pfmm_build_paras_at_level(l,scheme)
         END IF
         box_paras => mms_at_lev(l)%LHS_paras(:)
         stat_LHS_boxes(l) = SIZE(box_paras)
      CASE DEFAULT 
         CALL fmm_quit ('must select just LHS or RHS paras to use')
      END SELECT

   END SUBROUTINE pfmm_get_box_paras_at_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_get_box_qlm_at_level(l,scheme,qlm_T,side,memory)

      IMPLICIT NONE
      INTEGER(INTK),    INTENT(IN)  :: l
      TYPE(fmm_scheme), INTENT(IN)  :: scheme
      REAL(REALK),      POINTER     :: qlm_T(:,:)
      CHARACTER(3),     INTENT(IN)  :: side 
      CHARACTER(4),     INTENT(IN)  :: memory

      IF (.NOT.ASSOCIATED(mms_at_lev)) STOP 'mms_at_lev should be allocated!'
      IF (.NOT.ASSOCIATED(mms_at_lev(l)%qlm_T)) THEN 
         ! qlm data not available at this level yet. So make them!
         CALL pfmm_build_qlm_at_level(l,scheme,memory)
      END IF

      IF (side == 'LHS') CALL fmm_quit('currently no LHS boxed mms built!')
      IF (side == 'RHS') THEN
         qlm_T => mms_at_lev(l)%qlm_T(:,:)
      ELSE
         CALL fmm_quit("must select LHS or RHS boxed moments!")
      END IF

   END SUBROUTINE pfmm_get_box_qlm_at_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_preallocate_levels

      IMPLICIT NONE
      INTEGER(INTK) :: i

!FIXME add safety tests!!
      IF (pfmm_max_level == 0) &
         CALL fmm_quit('pfmm error allocating levels in box hierarchy')
!      IF (ASSOCIATED(mms_at_lev)) STOP 'mms_at_lev should not be allocated!'
!      IF (deepest_level < TOP_LEVEL) THEN
!         CALL fmm_quit('error allocating levels in box hierarchy')
!      END IF

      ! We allow for possibility of all levels being used, but sub-variables
      ! are only allocated if required, so this is not a big deal
      ALLOCATE (mms_at_lev(pfmm_max_level))
      DO i = LBOUND(mms_at_lev,1), UBOUND(mms_at_lev,1) 
         NULLIFY (mms_at_lev(i)%LHS_paras)
         NULLIFY (mms_at_lev(i)%RHS_paras)
         NULLIFY (mms_at_lev(i)%qlm_W)
         NULLIFY (mms_at_lev(i)%qlm_T)
      END DO

   END SUBROUTINE pfmm_preallocate_levels

!-------------------------------------------------------------------------------

   SUBROUTINE allocate_lm_at_level(l,mms_dim,LMAX)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: l, mms_dim, LMAX

      INTEGER(INTK) :: lm_dim
      LOGICAL :: fail

      lm_dim = (1+LMAX)**2
      IF (l > pfmm_max_level) CALL fmm_quit('pfmm invalid level to allocate!')
      IF (l < 1) CALL fmm_quit('pfmm invalid level to allocate!')

      ! must test if pointers already allocated (compiler may not notice)
      fail = .FALSE.
      IF (ASSOCIATED(mms_at_lev(l)%qlm_T)) fail = .TRUE.
      IF (ASSOCIATED(mms_at_lev(l)%qlm_W)) fail = .TRUE.
      IF (fail) CALL fmm_quit('box lm data already allocated!')
      ALLOCATE (mms_at_lev(l)%qlm_T(lm_dim,mms_dim))
      ALLOCATE (mms_at_lev(l)%qlm_W(lm_dim,mms_dim))
      ! must zero as they are built additively
      mms_at_lev(l)%qlm_T = zero
      mms_at_lev(l)%qlm_W = zero

   END SUBROUTINE allocate_lm_at_level

!-------------------------------------------------------------------------------

   SUBROUTINE pfmm_free_box_builder

      IMPLICIT NONE
      INTEGER(INTK) :: l

      NULLIFY (RHS_raw_mms)
!FIXME
      NULLIFY (LHS_raw_paras)

      IF (ASSOCIATED(mms_at_lev)) THEN
        DO l = LBOUND(mms_at_lev,1), UBOUND(mms_at_lev,1) 
          IF (ASSOCIATED(mms_at_lev(l)%LHS_paras,mms_at_lev(l)%RHS_paras)) THEN
             ! LHS and RHS paras the same and only point the same space
             DEALLOCATE(mms_at_lev(l)%RHS_paras)
          ELSE
             IF (ASSOCIATED(mms_at_lev(l)%RHS_paras))    &
                 DEALLOCATE(mms_at_lev(l)%RHS_paras)
             IF (ASSOCIATED(mms_at_lev(l)%LHS_paras))    &
                 DEALLOCATE(mms_at_lev(l)%LHS_paras)
          END IF
          IF (ASSOCIATED(mms_at_lev(l)%qlm_W)) DEALLOCATE(mms_at_lev(l)%qlm_W)
          IF (ASSOCIATED(mms_at_lev(l)%qlm_T)) DEALLOCATE(mms_at_lev(l)%qlm_T)
          NULLIFY (mms_at_lev(l)%LHS_paras)
          NULLIFY (mms_at_lev(l)%RHS_paras)
          NULLIFY (mms_at_lev(l)%qlm_W)
          NULLIFY (mms_at_lev(l)%qlm_T)
        END DO
        DEALLOCATE (mms_at_lev)
      END IF
      pfmm_max_level = 0

   END SUBROUTINE pfmm_free_box_builder

!-------------------------------------------------------------------------------

END MODULE pfmm_box_builder

